Skip to content

test indexing with arrays #341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 4 commits into from
Closed

test indexing with arrays #341

wants to merge 4 commits into from

Conversation

ev-br
Copy link
Member

@ev-br ev-br commented Feb 10, 2025 •
edited
Loading

cross-ref data-apis/array-api#669 : test indexing arrays objects with a mix of integers and arrays.

  • I'm not sure how to properly assert correctness here. I currently just check that results agree with numpy, which is not right. Maybe there's something in the ndindex library that would help. Could you weight in @asmeurer?
  • Do we need to test __setitem__ in addition to getitem?

As a data point, the current (hacky) test seems to pass locally with numpy (unsurprisingly), pytorch and jax.numpy.
EDIT: also passes locally with array-api-strict, patched via https://github.com/data-apis/array-api-strict/compare/main...ev-br:array_indexing?expand=1 (the patch certainly allows way too much, is for demo purposes only).

A CLI incantation to run tests locally 
$ ARRAY_API_TESTS_MODULE=torch pytest array_api_tests/test_array_object.py::test_getitem_arrays_and_ints -vs --max-examples=2000 --pdb

cc @kgryte for viz

@ev-br ev-br force-pushed the test_getitem_arrays branch from e4262e3 to bf9ed3f Compare February 10, 2025 16:58
@asmeurer
Copy link
Member

I think the main way to test correctness is to assert each element corresponds to indexing the individual element of the index array. This would make the test unvectorized, though. For a vectorized test the only thing I can think of is to flatten things and test that it matches take.

As far as mixed indices, ndindex can help by canonicalizing the index (you want the expand() method).

@ev-br
Copy link
Member Author

ev-br commented Feb 11, 2025

Run the (current, hacky) test on dask: it does not seem to allow array indexing.

$ ARRAY_API_TESTS_MODULE=array_api_compat.dask.array pytest array_api_tests/test_array_object.py::test_getitem_arrays_and_ints -vs --max-examples=2000 --pdb
...

??? x.shape = (1,) key = [dask.array<reshape, shape=(1, 1), dtype=int32, chunksize=(1, 1), chunktype=numpy.ndarray>] -- [(1, 1)]
FAILED
>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> traceback >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

array_api_tests/test_array_object.py:286: in test_getitem_arrays_and_ints
    out = x[key]
../../miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/core.py:2006: in __getitem__
    self, index2 = slice_with_int_dask_array(self, index2)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

x = dask.array<array, shape=(1,), dtype=int32, chunksize=(1,), chunktype=numpy.ndarray>
index = (dask.array<reshape, shape=(1, 1), dtype=int32, chunksize=(1, 1), chunktype=numpy.ndarray>,)

...

            if isinstance(idx, Array) and idx.dtype.kind in "iu":
                if idx.ndim == 0:
                    idx = idx[np.newaxis]
                    x = slice_with_int_dask_array_on_axis(x, idx, out_axis)
                    x = x[tuple(0 if i == out_axis else slice(None) for i in range(x.ndim))]
                    dropped_axis_cnt += 1
                elif idx.ndim == 1:
                    x = slice_with_int_dask_array_on_axis(x, idx, out_axis)
                    out_index.append(slice(None))
                else:
>                   raise NotImplementedError(
                        "Slicing with dask.array of ints only permitted when "
                        "the indexer has zero or one dimensions"
                    )
E                   NotImplementedError: Slicing with dask.array of ints only permitted when the indexer has zero or one dimensions
E                   Falsifying example: test_getitem_arrays_and_ints(
E                       shape=(1,),
E                       data=data(...),
E                   )
E                   Draw 1 (obj): [0]
E                   Draw 2: True
E                   Draw 3: ((1, 1),)
E                   Draw 4: dask.array<reshape, shape=(1, 1), dtype=int32, chunksize=(1, 1), chunktype=numpy.ndarray>
E                   Explanation:
E                       These lines were always and only run by failing examples:
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:344
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:346
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:348
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:349
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:350
E                           (and 4 more with settings.verbosity >= verbose)

../../miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/slicing.py:960: NotImplementedError

@ev-br ev-br force-pushed the test_getitem_arrays branch from bf9ed3f to ba9489a Compare February 11, 2025 16:56
@ev-br ev-br mentioned this pull request Feb 11, 2025
Closed
@ev-br
Copy link
Member Author

ev-br commented Feb 17, 2025

I think the main way to test correctness is to assert each element corresponds to indexing the individual element of the index array. This would make the test unvectorized, though. For a vectorized test the only thing I can think of is to flatten things and test that it matches take.

Thanks @asmeurer !
I actually wonder about a small generalization of what's currently in this PR: convert to numpy instead of an making it unvectorized. Will need a small utility to paper over downstream libs, but that's simple.
This is clearly a hack, even if a pragmatic one. I'm sure you have considered something like this previously @asmeurer ?

@kgryte kgryte mentioned this pull request Feb 17, 2025
Merged
@ev-br ev-br mentioned this pull request Feb 18, 2025
Merged
1 task
@ev-br
Copy link
Member Author

ev-br commented Feb 18, 2025

cross-ref #343 for a cleaner but slightly more restricted test (which only generates 1D indexing arrays)

@ev-br
Copy link
Member Author

ev-br commented Feb 18, 2025

closing in favor of #343

@ev-br ev-br closed this Feb 18, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@ev-br @asmeurer